#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Dec 16 01:17:53 2020

@author: agyeman
"""

###This is a vectorized version of the 1D Richards equation.
###The soil parameters at each node are the same.
##This code works for the flux and the zero gradient boundary conditions at the surface and the bottom of the field.

import numpy as np 
from Parameters_1D import * 
import van_Genuchten1D_vectorized as vg 


def Richards1D_vectorized(head, irrigAmount, cropCoeff, refEvap, soilPars):
    


    ##The spatial parameters
    totalDepth, axialNodes, axialDistance, totalNodes, numOfActuators, interPoints = spatialVariables_1D()

    capCapacity = vg.capillaryCapacity(head, soilPars) 

    ##Creating an array for the flux  at each node

    #flux = SX(np.zeros(totalNod1es+1))
    
    flux = np.zeros(totalNodes+1)

    ##Top boundary condition
    flux[totalNodes] = irrigAmount

    ##Bottom boundary
    #For the free drainage boundary condition, the flux equals the hydraulic conductivity of the bottom node
    flux[0] = -1*vg.hydraulicConductivity(head[0], soilPars)

    ##Flux for the internal nodes

    #The hydraulic conductivity at the center of the compartments
    cntr_1 = np.arange(0, totalNodes-1)
    nodalHydCond = vg.hydraulicConductivity(head, soilPars)
    approxHydCond = (nodalHydCond[cntr_1 + 1] + nodalHydCond[cntr_1])/2.0

    cntr_2 = np.arange(1,totalNodes)

    flux[cntr_2] = -approxHydCond*(((head[cntr_1 + 1]-head[cntr_1])/axialDistance) + 1.0)


    
    # sinkTerms = np.zeros((1, axialNodes))
    # sinkTerms = np.zeros(axialNodes)
    
    #dhdt = (-(flux[cntr_3 + 1] - flux[cntr_3])/ axialDistance)/capCapacity - ((cropCoeff*refEvap)/totalDepth)/capCapacity +  procDisturbance
    
    # i= np.arange(0, 1)
    distance = axialDistance*np.arange(0, axialNodes)
    sinkTerms = (2*(refEvap * cropCoeff)/ totalDepth)*(1-(totalDepth - distance)/totalDepth) #Weather conditions
    
    
    
    #dhdt = ((-(flux[cntr_3 + 1] - flux[cntr_3])/ axialDistance) - ((cropCoeff*refEvap)/totalDepth))/capCapacity
    
    cntr_3 = np.arange(0, totalNodes)
    
    dhdt = ((-(flux[cntr_3 + 1] - flux[cntr_3])/ axialDistance) - sinkTerms)/capCapacity
    return dhdt



def volMoistureAllNodes_1D(head, soilPars):
    
    totalDepth, axialNodes, axialDistance, totalNodes, numOfActuators, interPoints = spatialVariables_1D()
    volMoistureInit = vg.volumetricMoisture(head, soilPars)
    cMatrix = CMatrixAllMeasurements(totalNodes)
    volMoistureFin = np.dot(cMatrix, volMoistureInit)
    return volMoistureFin


##For this model,the cMatrix is defined as follows

def cMatrix_1D(axialNodes):

    cMatrix = np.zeros(axialNodes)

    for i in range(axialNodes):
        cMatrix[i] = (1/axialNodes)
        
    cMatrix = np.reshape(cMatrix, (1,16))

    return cMatrix



def volMoistureSelected_1D(cMatrix, head, soilPars):
    volMoistureInit = vg.volumetricMoisture(head, soilPars)
    volMoistureFin = np.dot(cMatrix, volMoistureInit)
    return volMoistureFin